{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "395e5929",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "#全局变量\n",
    "hub_token = open('/root/hub_token.txt').read().strip()\n",
    "repo_id = 'lansinuote/diffusion.8.instruct_pix2pix'\n",
    "push_to_hub = True\n",
    "checkpoint = 'runwayml/stable-diffusion-v1-5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9d841e29",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration lansinuote--diffusion.8.instruct_pix2pix-5db27ce94b1e4a9e\n",
      "Found cached dataset parquet (/root/.cache/huggingface/datasets/lansinuote___parquet/lansinuote--diffusion.8.instruct_pix2pix-5db27ce94b1e4a9e/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(Dataset({\n",
       "     features: ['input', 'text', 'output'],\n",
       "     num_rows: 1000\n",
       " }),\n",
       " {'input': tensor([[[-0.3882, -0.3647, -0.3569,  ...,  0.5529,  0.2235,  0.0745],\n",
       "           [-0.3804, -0.3569, -0.3490,  ...,  0.4118,  0.0510, -0.1373],\n",
       "           [-0.3804, -0.3647, -0.3333,  ...,  0.0902, -0.1451, -0.2941],\n",
       "           ...,\n",
       "           [ 0.7647,  0.7569,  0.7490,  ..., -0.1451, -0.1765, -0.1529],\n",
       "           [ 0.7412,  0.7333,  0.7569,  ..., -0.1451, -0.1686, -0.1686],\n",
       "           [ 0.7412,  0.7255,  0.7412,  ..., -0.1373, -0.1451, -0.1765]],\n",
       "  \n",
       "          [[-0.8431, -0.8196, -0.8118,  ...,  0.2392, -0.1059, -0.2627],\n",
       "           [-0.8275, -0.8196, -0.8118,  ...,  0.0510, -0.2941, -0.4667],\n",
       "           [-0.8118, -0.8196, -0.7961,  ..., -0.2941, -0.4745, -0.5529],\n",
       "           ...,\n",
       "           [ 0.6392,  0.6314,  0.6235,  ..., -0.5843, -0.6000, -0.5765],\n",
       "           [ 0.6157,  0.6078,  0.6314,  ..., -0.5843, -0.6000, -0.6000],\n",
       "           [ 0.6157,  0.6078,  0.6235,  ..., -0.5765, -0.5843, -0.6078]],\n",
       "  \n",
       "          [[-0.8510, -0.8353, -0.8275,  ..., -0.2549, -0.5451, -0.6157],\n",
       "           [-0.8510, -0.8431, -0.8275,  ..., -0.3490, -0.6314, -0.7255],\n",
       "           [-0.8353, -0.8510, -0.8118,  ..., -0.5686, -0.6863, -0.7255],\n",
       "           ...,\n",
       "           [ 0.4980,  0.4980,  0.4980,  ..., -0.7647, -0.7882, -0.7647],\n",
       "           [ 0.4745,  0.4667,  0.4980,  ..., -0.7804, -0.7961, -0.7804],\n",
       "           [ 0.4824,  0.4588,  0.4745,  ..., -0.7569, -0.7647, -0.7804]]]),\n",
       "  'output': tensor([[[-0.4353, -0.4275, -0.4353,  ...,  0.4039,  0.3725,  0.3882],\n",
       "           [-0.4353, -0.4431, -0.4353,  ...,  0.5059,  0.4275,  0.3647],\n",
       "           [-0.4667, -0.4431, -0.4118,  ...,  0.4902,  0.4824,  0.5059],\n",
       "           ...,\n",
       "           [ 0.7490,  0.7569,  0.7569,  ..., -0.0667, -0.0902, -0.0824],\n",
       "           [ 0.7412,  0.7569,  0.7569,  ..., -0.0588, -0.0902, -0.0824],\n",
       "           [ 0.7412,  0.7490,  0.7490,  ..., -0.0510, -0.0588, -0.0667]],\n",
       "  \n",
       "          [[-0.5059, -0.4980, -0.4980,  ...,  0.4196,  0.3804,  0.3961],\n",
       "           [-0.4902, -0.5137, -0.4902,  ...,  0.5059,  0.4275,  0.3569],\n",
       "           [-0.5216, -0.5059, -0.4745,  ...,  0.4902,  0.4745,  0.4980],\n",
       "           ...,\n",
       "           [ 0.7176,  0.7255,  0.7255,  ..., -0.1216, -0.1373, -0.1294],\n",
       "           [ 0.7098,  0.7255,  0.7255,  ..., -0.1137, -0.1529, -0.1451],\n",
       "           [ 0.7098,  0.7176,  0.7176,  ..., -0.1137, -0.1216, -0.1373]],\n",
       "  \n",
       "          [[-0.4824, -0.4745, -0.4745,  ...,  0.3804,  0.3333,  0.3569],\n",
       "           [-0.4902, -0.4980, -0.4745,  ...,  0.4824,  0.4039,  0.3333],\n",
       "           [-0.5137, -0.4902, -0.4510,  ...,  0.4667,  0.4588,  0.4824],\n",
       "           ...,\n",
       "           [ 0.7020,  0.7020,  0.7098,  ..., -0.1451, -0.1608, -0.1529],\n",
       "           [ 0.6863,  0.7020,  0.7098,  ..., -0.1294, -0.1686, -0.1608],\n",
       "           [ 0.6863,  0.6941,  0.7098,  ..., -0.1294, -0.1451, -0.1608]]]),\n",
       "  'text': tensor([49406,  2105,   585,  1095,   320,  1125, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407])})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from transformers import CLIPTokenizer\n",
    "import torchvision\n",
    "\n",
    "tokenizer = CLIPTokenizer.from_pretrained(checkpoint, subfolder='tokenizer')\n",
    "\n",
    "compose = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.Resize(256),\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    lambda x: (x * 2) - 1,\n",
    "])\n",
    "\n",
    "#使用我转存的数据集\n",
    "#fusing/instructpix2pix-1000-samples\n",
    "dataset = load_dataset(path=repo_id, split='train')\n",
    "\n",
    "\n",
    "def f(data):\n",
    "    #图像编码\n",
    "    input = [compose(i) for i in data['input']]\n",
    "    output = [compose(i) for i in data['output']]\n",
    "\n",
    "    #文字编码\n",
    "    #77 = tokenizer.model_max_length\n",
    "    text = tokenizer.batch_encode_plus(data['text'],\n",
    "                                       max_length=77,\n",
    "                                       padding='max_length',\n",
    "                                       truncation=True,\n",
    "                                       return_tensors='pt').input_ids\n",
    "\n",
    "    return {'input': input, 'output': output, 'text': text}\n",
    "\n",
    "\n",
    "dataset = dataset.with_transform(f)\n",
    "\n",
    "dataset, dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "43037e7c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(250,\n",
       " {'input': tensor([[[[-0.0667, -0.0824, -0.1059,  ...,  0.4588,  0.4667,  0.5059],\n",
       "            [ 0.1922,  0.1686,  0.1529,  ...,  0.4118,  0.4510,  0.4745],\n",
       "            [ 0.4902,  0.5137,  0.5216,  ...,  0.4039,  0.3882,  0.4510],\n",
       "            ...,\n",
       "            [ 0.2627,  0.2471,  0.2235,  ..., -0.9294, -0.9373, -0.9843],\n",
       "            [ 0.2784,  0.2392,  0.1137,  ..., -0.9137, -0.9373, -0.9608],\n",
       "            [ 0.2471,  0.1373,  0.1529,  ..., -0.9059, -0.9294, -0.9059]],\n",
       "  \n",
       "           [[-0.3647, -0.3647, -0.3882,  ...,  0.4275,  0.4275,  0.4745],\n",
       "            [-0.1451, -0.1529, -0.1765,  ...,  0.3804,  0.4196,  0.4353],\n",
       "            [ 0.1373,  0.1529,  0.1608,  ...,  0.3882,  0.3804,  0.4275],\n",
       "            ...,\n",
       "            [ 0.0275,  0.0196, -0.0196,  ..., -0.1373, -0.2000, -0.2863],\n",
       "            [ 0.0275,  0.0039, -0.1137,  ..., -0.0588, -0.1216, -0.1922],\n",
       "            [ 0.0118, -0.1059, -0.0824,  ..., -0.0275, -0.0824, -0.1294]],\n",
       "  \n",
       "           [[-0.7961, -0.7882, -0.7961,  ...,  0.2000,  0.2000,  0.2471],\n",
       "            [-0.7020, -0.7020, -0.7176,  ...,  0.1529,  0.1843,  0.2078],\n",
       "            [-0.5765, -0.5529, -0.5529,  ...,  0.1686,  0.1529,  0.2000],\n",
       "            ...,\n",
       "            [ 0.3961,  0.3804,  0.3490,  ...,  0.9137,  0.8902,  0.8431],\n",
       "            [ 0.3961,  0.3725,  0.2549,  ...,  0.9294,  0.9137,  0.8980],\n",
       "            [ 0.3647,  0.2627,  0.2863,  ...,  0.9373,  0.9137,  0.8902]]],\n",
       "  \n",
       "  \n",
       "          [[[-0.5529, -0.5373, -0.4980,  ...,  0.3020,  0.1529, -0.2863],\n",
       "            [-0.8196, -0.7804, -0.7490,  ...,  0.0902, -0.1686, -0.4902],\n",
       "            [-0.9059, -0.8902, -0.8745,  ..., -0.0196, -0.3255, -0.5922],\n",
       "            ...,\n",
       "            [-0.5373, -0.4431, -0.3647,  ...,  0.4039,  0.3882,  0.3647],\n",
       "            [-0.6235, -0.5529, -0.4824,  ...,  0.3020,  0.3020,  0.2863],\n",
       "            [-0.6941, -0.6392, -0.6078,  ...,  0.2784,  0.2314,  0.2000]],\n",
       "  \n",
       "           [[-0.2078, -0.1922, -0.1451,  ...,  0.7961,  0.5843,  0.0902],\n",
       "            [-0.4118, -0.3725, -0.3333,  ...,  0.6392,  0.3098, -0.0902],\n",
       "            [-0.4980, -0.4667, -0.4275,  ...,  0.5608,  0.1608, -0.2157],\n",
       "            ...,\n",
       "            [-0.9373, -0.9137, -0.8980,  ..., -0.4745, -0.4745, -0.4824],\n",
       "            [-0.9451, -0.9216, -0.9059,  ..., -0.5451, -0.5451, -0.5529],\n",
       "            [-0.9529, -0.9373, -0.9294,  ..., -0.5608, -0.5922, -0.5922]],\n",
       "  \n",
       "           [[ 0.1451,  0.1608,  0.1922,  ...,  0.8275,  0.7020,  0.2706],\n",
       "            [ 0.0275,  0.0667,  0.1059,  ...,  0.7255,  0.4824,  0.1451],\n",
       "            [-0.0039,  0.0588,  0.0980,  ...,  0.6706,  0.3725,  0.0588],\n",
       "            ...,\n",
       "            [-0.9765, -0.9765, -0.9765,  ..., -0.9765, -0.9451, -0.9216],\n",
       "            [-0.9686, -0.9686, -0.9608,  ..., -0.9686, -0.9529, -0.9451],\n",
       "            [-0.9686, -0.9686, -0.9765,  ..., -0.9451, -0.9608, -0.9373]]],\n",
       "  \n",
       "  \n",
       "          [[[-0.5059, -0.5765, -0.5137,  ...,  0.3725, -0.0902, -0.7490],\n",
       "            [-0.4824, -0.5765, -0.4980,  ...,  0.2392, -0.1294, -0.7412],\n",
       "            [-0.4353, -0.5451, -0.5059,  ...,  0.1529, -0.2000, -0.7333],\n",
       "            ...,\n",
       "            [-0.0745, -0.0824, -0.6157,  ..., -0.9765, -0.8118, -0.2627],\n",
       "            [-0.1137, -0.1373, -0.5843,  ..., -0.9765, -0.7882, -0.2235],\n",
       "            [-0.0902, -0.0431, -0.5451,  ..., -0.9843, -0.7882, -0.2078]],\n",
       "  \n",
       "           [[-0.4118, -0.3882, -0.3333,  ...,  0.3098, -0.1608, -0.8039],\n",
       "            [-0.3725, -0.3961, -0.3333,  ...,  0.1765, -0.2078, -0.7961],\n",
       "            [-0.3647, -0.3882, -0.3412,  ...,  0.1059, -0.2706, -0.8039],\n",
       "            ...,\n",
       "            [-0.2314, -0.2235, -0.6706,  ..., -0.9608, -0.8039, -0.2784],\n",
       "            [-0.2784, -0.2706, -0.6314,  ..., -0.9608, -0.7882, -0.2471],\n",
       "            [-0.2627, -0.1765, -0.5922,  ..., -0.9686, -0.7882, -0.2392]],\n",
       "  \n",
       "           [[-0.8196, -0.8275, -0.8196,  ...,  0.0039, -0.3647, -0.8980],\n",
       "            [-0.8275, -0.8510, -0.8196,  ..., -0.0902, -0.4118, -0.9059],\n",
       "            [-0.8118, -0.8275, -0.8275,  ..., -0.1843, -0.4980, -0.9294],\n",
       "            ...,\n",
       "            [-0.6549, -0.5843, -0.8196,  ..., -0.9843, -0.9529, -0.6627],\n",
       "            [-0.7098, -0.6314, -0.7882,  ..., -0.9843, -0.9373, -0.6235],\n",
       "            [-0.6941, -0.5608, -0.7725,  ..., -0.9765, -0.9216, -0.5843]]],\n",
       "  \n",
       "  \n",
       "          [[[-0.6471, -0.6784, -0.7020,  ..., -0.5137, -0.5294, -0.5059],\n",
       "            [-0.6706, -0.7176, -0.7098,  ..., -0.5608, -0.5765, -0.5137],\n",
       "            [-0.6706, -0.7098, -0.7176,  ..., -0.5843, -0.5922, -0.5059],\n",
       "            ...,\n",
       "            [ 0.6549,  0.7020,  0.7255,  ...,  0.4588,  0.4510,  0.4431],\n",
       "            [ 0.6941,  0.7020,  0.7569,  ...,  0.5373,  0.5059,  0.5373],\n",
       "            [ 0.8039,  0.8039,  0.8431,  ...,  0.6627,  0.6392,  0.6549]],\n",
       "  \n",
       "           [[-0.4667, -0.4745, -0.4824,  ..., -0.3412, -0.3569, -0.3412],\n",
       "            [-0.4745, -0.5059, -0.4902,  ..., -0.3647, -0.3882, -0.3412],\n",
       "            [-0.4902, -0.5059, -0.4902,  ..., -0.3804, -0.3961, -0.3333],\n",
       "            ...,\n",
       "            [ 0.7882,  0.8275,  0.8431,  ...,  0.6784,  0.6627,  0.6471],\n",
       "            [ 0.7961,  0.8039,  0.8588,  ...,  0.7255,  0.6863,  0.7020],\n",
       "            [ 0.8824,  0.8902,  0.9216,  ...,  0.8118,  0.7882,  0.7882]],\n",
       "  \n",
       "           [[-0.7804, -0.8275, -0.8510,  ..., -0.6392, -0.6627, -0.6314],\n",
       "            [-0.8039, -0.8667, -0.8667,  ..., -0.6784, -0.6941, -0.6471],\n",
       "            [-0.8196, -0.8510, -0.8588,  ..., -0.6941, -0.7020, -0.6314],\n",
       "            ...,\n",
       "            [ 0.8196,  0.8431,  0.8510,  ...,  0.8118,  0.7961,  0.7725],\n",
       "            [ 0.8196,  0.8118,  0.8588,  ...,  0.8196,  0.7804,  0.7882],\n",
       "            [ 0.8980,  0.8980,  0.9216,  ...,  0.8510,  0.8353,  0.8275]]]]),\n",
       "  'output': tensor([[[[ 0.3255,  0.3020,  0.2549,  ...,  0.8118,  0.9765,  0.9843],\n",
       "            [ 0.4588,  0.4275,  0.3882,  ...,  0.8196,  0.9843,  0.9843],\n",
       "            [ 0.5765,  0.5294,  0.5059,  ...,  0.8275,  0.9843,  0.9765],\n",
       "            ...,\n",
       "            [ 0.6157,  0.6157,  0.5922,  ..., -0.6314, -0.7176, -0.7882],\n",
       "            [ 0.5373,  0.6000,  0.6078,  ..., -0.6392, -0.7412, -0.8118],\n",
       "            [ 0.4745,  0.5137,  0.5843,  ..., -0.6627, -0.7490, -0.7725]],\n",
       "  \n",
       "           [[-0.3255, -0.3647, -0.4196,  ...,  0.7569,  0.9765,  0.9922],\n",
       "            [-0.2471, -0.2627, -0.3176,  ...,  0.7647,  0.9765,  0.9843],\n",
       "            [-0.1294, -0.1765, -0.1922,  ...,  0.7725,  0.9765,  0.9843],\n",
       "            ...,\n",
       "            [ 0.7412,  0.7569,  0.7333,  ...,  0.1216,  0.0118, -0.0902],\n",
       "            [ 0.6627,  0.7412,  0.7412,  ...,  0.1294, -0.0039, -0.1137],\n",
       "            [ 0.6078,  0.6549,  0.7176,  ...,  0.0980, -0.0118, -0.1059]],\n",
       "  \n",
       "           [[-0.7255, -0.7569, -0.7725,  ...,  0.4588,  0.8196,  0.8902],\n",
       "            [-0.7020, -0.7098, -0.7255,  ...,  0.4667,  0.8196,  0.8902],\n",
       "            [-0.6471, -0.6863, -0.7020,  ...,  0.4745,  0.8275,  0.8902],\n",
       "            ...,\n",
       "            [ 1.0000,  1.0000,  0.9843,  ...,  1.0000,  0.9843,  0.9765],\n",
       "            [ 0.9765,  1.0000,  0.9922,  ...,  0.9922,  0.9765,  0.9451],\n",
       "            [ 0.9529,  0.9686,  0.9922,  ...,  0.9843,  0.9686,  0.9059]]],\n",
       "  \n",
       "  \n",
       "          [[[-0.5608, -0.5451, -0.5373,  ...,  0.6157,  0.4275,  0.0745],\n",
       "            [-0.6863, -0.6863, -0.6784,  ...,  0.5137,  0.2471, -0.0510],\n",
       "            [-0.7490, -0.7569, -0.7412,  ...,  0.4275,  0.0902, -0.1451],\n",
       "            ...,\n",
       "            [-0.7020, -0.6784, -0.6549,  ...,  0.0196,  0.0196, -0.0275],\n",
       "            [-0.7569, -0.7490, -0.7333,  ..., -0.0275, -0.0510, -0.0824],\n",
       "            [-0.8039, -0.8039, -0.7961,  ...,  0.0118, -0.0588, -0.1137]],\n",
       "  \n",
       "           [[-0.3804, -0.3569, -0.3569,  ...,  0.8588,  0.6549,  0.2941],\n",
       "            [-0.4745, -0.4824, -0.4745,  ...,  0.7725,  0.4902,  0.1843],\n",
       "            [-0.5373, -0.5373, -0.5216,  ...,  0.7020,  0.3490,  0.0824],\n",
       "            ...,\n",
       "            [-0.9608, -0.9608, -0.9451,  ..., -0.5922, -0.5608, -0.5843],\n",
       "            [-0.9608, -0.9686, -0.9529,  ..., -0.6235, -0.6314, -0.6471],\n",
       "            [-0.9686, -0.9765, -0.9608,  ..., -0.5843, -0.6314, -0.6627]],\n",
       "  \n",
       "           [[-0.3176, -0.3020, -0.2941,  ...,  0.8745,  0.7020,  0.3647],\n",
       "            [-0.3961, -0.4039, -0.4039,  ...,  0.8118,  0.5608,  0.2627],\n",
       "            [-0.4431, -0.4431, -0.4353,  ...,  0.7490,  0.4196,  0.1608],\n",
       "            ...,\n",
       "            [-0.9922, -0.9922, -0.9843,  ..., -0.8980, -0.8510, -0.8353],\n",
       "            [-0.9765, -0.9843, -0.9686,  ..., -0.8980, -0.8824, -0.8824],\n",
       "            [-0.9765, -0.9843, -0.9686,  ..., -0.8353, -0.8588, -0.8667]]],\n",
       "  \n",
       "  \n",
       "          [[[-0.3176, -0.5216, -0.4902,  ..., -0.4353, -0.4824, -0.6314],\n",
       "            [-0.3255, -0.5059, -0.4745,  ..., -0.4039, -0.4745, -0.6078],\n",
       "            [-0.3098, -0.4667, -0.4667,  ..., -0.4039, -0.5137, -0.5608],\n",
       "            ...,\n",
       "            [-0.7412, -0.7961, -0.8667,  ..., -0.8902, -0.8745, -0.7647],\n",
       "            [-0.7412, -0.7490, -0.8431,  ..., -0.8902, -0.8824, -0.7725],\n",
       "            [-0.7412, -0.7490, -0.8196,  ..., -0.8745, -0.8588, -0.8039]],\n",
       "  \n",
       "           [[-0.3804, -0.5529, -0.5216,  ..., -0.5137, -0.5529, -0.6941],\n",
       "            [-0.3725, -0.5373, -0.5059,  ..., -0.4745, -0.5451, -0.6784],\n",
       "            [-0.3647, -0.5059, -0.5137,  ..., -0.4745, -0.5922, -0.6471],\n",
       "            ...,\n",
       "            [-0.7255, -0.7725, -0.8510,  ..., -0.8588, -0.8510, -0.7490],\n",
       "            [-0.7255, -0.7255, -0.8118,  ..., -0.8588, -0.8510, -0.7490],\n",
       "            [-0.7333, -0.7255, -0.8039,  ..., -0.8510, -0.8431, -0.7882]],\n",
       "  \n",
       "           [[-0.6863, -0.8353, -0.7961,  ..., -0.7882, -0.8118, -0.9216],\n",
       "            [-0.6941, -0.8118, -0.7725,  ..., -0.7647, -0.8039, -0.9216],\n",
       "            [-0.6941, -0.7725, -0.7804,  ..., -0.7647, -0.8431, -0.8824],\n",
       "            ...,\n",
       "            [-0.7725, -0.8196, -0.8745,  ..., -0.9059, -0.8980, -0.8196],\n",
       "            [-0.7647, -0.7725, -0.8431,  ..., -0.8980, -0.9059, -0.8196],\n",
       "            [-0.7725, -0.7647, -0.8275,  ..., -0.8902, -0.8980, -0.8588]]],\n",
       "  \n",
       "  \n",
       "          [[[-0.6627, -0.7255, -0.7333,  ..., -0.6392, -0.6549, -0.6000],\n",
       "            [-0.7020, -0.7412, -0.7255,  ..., -0.6627, -0.6706, -0.6314],\n",
       "            [-0.6706, -0.7255, -0.7176,  ..., -0.6706, -0.6941, -0.6314],\n",
       "            ...,\n",
       "            [ 0.2078,  0.2314,  0.2471,  ..., -0.0431, -0.0667, -0.0745],\n",
       "            [ 0.2078,  0.2235,  0.2784,  ...,  0.0275, -0.0039, -0.0039],\n",
       "            [ 0.2863,  0.3020,  0.3490,  ...,  0.1608,  0.1216,  0.1294]],\n",
       "  \n",
       "           [[ 0.1608,  0.1451,  0.1529,  ...,  0.1137,  0.0824,  0.1059],\n",
       "            [ 0.1608,  0.1373,  0.1529,  ...,  0.1137,  0.1059,  0.1059],\n",
       "            [ 0.1686,  0.1373,  0.1451,  ...,  0.1059,  0.0745,  0.0980],\n",
       "            ...,\n",
       "            [ 0.8118,  0.8275,  0.8275,  ...,  0.6471,  0.6392,  0.6235],\n",
       "            [ 0.8196,  0.8118,  0.8353,  ...,  0.6706,  0.6627,  0.6549],\n",
       "            [ 0.8745,  0.8667,  0.8745,  ...,  0.7412,  0.7255,  0.7098]],\n",
       "  \n",
       "           [[-0.0824, -0.1216, -0.1294,  ..., -0.1216, -0.1373, -0.0667],\n",
       "            [-0.0980, -0.1451, -0.1216,  ..., -0.1294, -0.1216, -0.0902],\n",
       "            [-0.0824, -0.1373, -0.1294,  ..., -0.1373, -0.1529, -0.0902],\n",
       "            ...,\n",
       "            [ 0.4824,  0.4980,  0.4824,  ...,  0.4196,  0.4196,  0.4039],\n",
       "            [ 0.4745,  0.4588,  0.4902,  ...,  0.4118,  0.4118,  0.4275],\n",
       "            [ 0.5451,  0.5294,  0.5529,  ...,  0.4588,  0.4588,  0.4745]]]]),\n",
       "  'text': tensor([[49406,  1078,   585,   320, 27014, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407],\n",
       "          [49406, 15234,   620,   518, 30963,  7244,   593,   320, 20075,  2390,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407],\n",
       "          [49406,   720,  1269,  3309,   550,  5034,  8721, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407],\n",
       "          [49406,  1078,  1180, 20387, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407]])})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loader = torch.utils.data.DataLoader(dataset,\n",
    "                                     shuffle=True,\n",
    "                                     collate_fn=None,\n",
    "                                     batch_size=4)\n",
    "\n",
    "len(loader), next(iter(loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "217c0a9c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "encoder 12306.048\n",
      "vae 8365.3863\n",
      "unet 85953.2484\n"
     ]
    }
   ],
   "source": [
    "from diffusers import AutoencoderKL, UNet2DConditionModel, ControlNetModel\n",
    "from transformers import CLIPTextModel\n",
    "\n",
    "#加载3个模型\n",
    "encoder = CLIPTextModel.from_pretrained(checkpoint, subfolder='text_encoder')\n",
    "vae = AutoencoderKL.from_pretrained(checkpoint, subfolder='vae')\n",
    "unet = UNet2DConditionModel.from_pretrained(checkpoint, subfolder='unet')\n",
    "\n",
    "# #修改unet.conv_in层的形状\n",
    "unet.register_to_config(in_channels=8)\n",
    "with torch.no_grad():\n",
    "    new_conv_in = torch.nn.Conv2d(8, 320, 3, 1, 1)\n",
    "    new_conv_in.weight.zero_()\n",
    "    new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)\n",
    "    unet.conv_in = new_conv_in\n",
    "print(unet.config.in_channels)\n",
    "\n",
    "# encoder = CLIPTextModel.from_pretrained('./save', subfolder='text_encoder')\n",
    "# vae = AutoencoderKL.from_pretrained('./save', subfolder='vae')\n",
    "# unet = UNet2DConditionModel.from_pretrained('./save', subfolder='unet')\n",
    "\n",
    "\n",
    "def print_model_size(name, model):\n",
    "    print(name, sum(i.numel() for i in model.parameters()) / 10000)\n",
    "\n",
    "\n",
    "print_model_size('encoder', encoder)\n",
    "print_model_size('vae', vae)\n",
    "print_model_size('unet', unet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a5f838ac",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DDPMScheduler {\n",
       "   \"_class_name\": \"DDPMScheduler\",\n",
       "   \"_diffusers_version\": \"0.15.1\",\n",
       "   \"beta_end\": 0.012,\n",
       "   \"beta_schedule\": \"scaled_linear\",\n",
       "   \"beta_start\": 0.00085,\n",
       "   \"clip_sample\": false,\n",
       "   \"clip_sample_range\": 1.0,\n",
       "   \"dynamic_thresholding_ratio\": 0.995,\n",
       "   \"num_train_timesteps\": 1000,\n",
       "   \"prediction_type\": \"epsilon\",\n",
       "   \"sample_max_value\": 1.0,\n",
       "   \"set_alpha_to_one\": false,\n",
       "   \"skip_prk_steps\": true,\n",
       "   \"steps_offset\": 1,\n",
       "   \"thresholding\": false,\n",
       "   \"trained_betas\": null,\n",
       "   \"variance_type\": \"fixed_small\"\n",
       " },\n",
       " AdamW (\n",
       " Parameter Group 0\n",
       "     amsgrad: False\n",
       "     betas: (0.9, 0.999)\n",
       "     capturable: False\n",
       "     eps: 1e-08\n",
       "     foreach: None\n",
       "     lr: 5e-05\n",
       "     maximize: False\n",
       "     weight_decay: 0.01\n",
       " ),\n",
       " MSELoss())"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from diffusers import DDPMScheduler\n",
    "\n",
    "scheduler = DDPMScheduler.from_pretrained(checkpoint, subfolder='scheduler')\n",
    "\n",
    "optimizer = torch.optim.AdamW(unet.parameters(),\n",
    "                              lr=5e-5,\n",
    "                              betas=(0.9, 0.999),\n",
    "                              weight_decay=0.01,\n",
    "                              eps=1e-8)\n",
    "\n",
    "criterion = torch.nn.MSELoss()\n",
    "\n",
    "scheduler, optimizer, criterion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b766d5c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dropout_data(out_encoder, input):\n",
    "    #输入图编码\n",
    "    #[4, 3, 256, 256] -> [4, 4, 32, 32]\n",
    "    out_vae_input = vae.encode(input).latent_dist.mode()\n",
    "\n",
    "    #做mask\n",
    "    r = torch.rand(4, device=out_encoder.device)\n",
    "    #[4, 1, 1]\n",
    "    mask_text = (r > 0.1).reshape(4, 1, 1)\n",
    "    #[4, 1, 1, 1]\n",
    "    mask_image = torch.logical_or(r < 0.05,\n",
    "                                  r > 0.15).float().reshape(4, 1, 1, 1)\n",
    "    \n",
    "    #编码负采样的文本,其实这里只要算一次就行了\n",
    "    out_encoder_neg = tokenizer.batch_encode_plus(\n",
    "        [''],\n",
    "        max_length=77,\n",
    "        padding='max_length',\n",
    "        truncation=True,\n",
    "        return_tensors='pt').input_ids.to(out_encoder.device)\n",
    "    out_encoder_neg = encoder(out_encoder_neg)[0]\n",
    "\n",
    "    #使用mask混合正负编码\n",
    "    #文本大概率选择正编码\n",
    "    #[4, 77, 768]\n",
    "    out_encoder = torch.where(mask_text, out_encoder, out_encoder_neg)\n",
    "\n",
    "    #图像小概率归零\n",
    "    #[4, 4, 32, 32]\n",
    "    out_vae_input = mask_image * out_vae_input\n",
    "\n",
    "    return out_encoder, out_vae_input\n",
    "\n",
    "# dropout_data(out_encoder=torch.randn(4, 77, 768),\n",
    "#              input=torch.randn(4, 3, 256, 256))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a3f32e9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_loss(data):\n",
    "    #文字编码\n",
    "    #[4, 77] -> [4, 77, 768]\n",
    "    out_encoder = encoder(data['text'])[0]\n",
    "\n",
    "    #输出图编码\n",
    "    #[4, 3, 256, 256] -> [4, 4, 32, 32]\n",
    "    out_vae_output = vae.encode(data['output']).latent_dist.sample()\n",
    "    #0.18215 = vae.config.scaling_factor\n",
    "    out_vae_output = out_vae_output * 0.18215\n",
    "\n",
    "    #随机数,unet的计算目标\n",
    "    #[4, 4, 32, 32]\n",
    "    noise = torch.randn_like(out_vae_output)\n",
    "\n",
    "    #往特征图中添加噪声\n",
    "    #1000 = scheduler.num_train_timesteps\n",
    "    #4 = out_vae.shape[0]\n",
    "    noise_step = torch.randint(0, 1000, (4, )).long()\n",
    "    noise_step = noise_step.to(out_encoder.device)\n",
    "    #[4, 4, 32, 32]\n",
    "    out_vae_noise = scheduler.add_noise(out_vae_output, noise, noise_step)\n",
    "\n",
    "    #使用mask组合正负采样的文本编码数据\n",
    "    #输入图编码\n",
    "    #[4, 77, 768],[4, 4, 32, 32]\n",
    "    out_encoder, out_vae_input = dropout_data(out_encoder=out_encoder,\n",
    "                                              input=data['input'])\n",
    "\n",
    "    #向out_vae_noise中组合输入图的数据\n",
    "    #[4, 4+4, 32, 32] -> [4, 8, 32, 32]\n",
    "    out_vae_noise = torch.cat([out_vae_noise, out_vae_input], dim=1)\n",
    "\n",
    "    #根据文字信息,把特征图中的噪声计算出来\n",
    "    #[4, 4, 32, 32]\n",
    "    out_unet = unet(out_vae_noise,\n",
    "                    noise_step,\n",
    "                    encoder_hidden_states=out_encoder).sample\n",
    "\n",
    "    #计算mse loss\n",
    "    #[4, 4, 32, 32],[4, 4, 32, 32]\n",
    "    return criterion(out_unet, noise)\n",
    "\n",
    "\n",
    "# get_loss({\n",
    "#     'text': torch.ones(4, 77).long(),\n",
    "#     'input': torch.randn(4, 3, 256, 256),\n",
    "#     'output': torch.randn(4, 3, 256, 256),\n",
    "# })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "84b88599",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from diffusers import StableDiffusionPipeline\n",
    "from huggingface_hub import Repository, create_repo\n",
    "\n",
    "\n",
    "def train():\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "    unet.to(device)\n",
    "    encoder.to(device)\n",
    "    vae.to(device)\n",
    "\n",
    "    vae.requires_grad_(False)\n",
    "    encoder.requires_grad_(False)\n",
    "    unet.train()\n",
    "\n",
    "    loss_sum = 0\n",
    "    for epoch in range(4000):\n",
    "        for i, data in enumerate(loader):\n",
    "            for k in data.keys():\n",
    "                data[k] = data[k].to(device)\n",
    "\n",
    "            loss = get_loss(data) / 4\n",
    "            loss.backward()\n",
    "            loss_sum += loss.item()\n",
    "\n",
    "            if i % 4 == 0:\n",
    "                torch.nn.utils.clip_grad_norm_(unet.parameters(), 1.0)\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "        if (epoch + 1) % 20 == 0:\n",
    "            print(epoch, loss_sum)\n",
    "            loss_sum = 0\n",
    "\n",
    "            #保存\n",
    "            StableDiffusionPipeline.from_pretrained(\n",
    "                checkpoint, text_encoder=encoder, vae=vae,\n",
    "                unet=unet).save_pretrained('./save')\n",
    "\n",
    "\n",
    "if push_to_hub:\n",
    "    create_repo(repo_id, exist_ok=True, token=hub_token)\n",
    "    repo = Repository('./save', clone_from=repo_id, token=hub_token)\n",
    "    train()\n",
    "    repo.push_to_hub()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pt39]",
   "language": "python",
   "name": "conda-env-pt39-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
